package com.hsl.housaileibot001.ai.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ResourceLoader;

import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 * @author liul
 * @date 2025/10/18 20:41
 */
// @Configuration
public class RagStandardConfig {

    // ! 这里注入的是配置文件中的 text-embedding-v4 模型（1024维），如果没有配置文件，默认使用内部的模型（384维）
    @Resource
    private EmbeddingModel qwenEmbeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Resource
    private ResourceLoader resourceLoader;

    @Bean
    public ContentRetriever contentRetriever() {
        // 1. 加载文档（todo:这里使用内存存储文档，后续使用向量数据库！）
        // 支持多种路径：优先使用 classpath，其次是文件系统路径
        String docsPath = null;
        try {
            // 尝试从 classpath 加载
            org.springframework.core.io.Resource resource = resourceLoader.getResource("classpath:docs");
            if (resource.exists() && resource.getFile().isDirectory()) {
                docsPath = resource.getFile().getAbsolutePath();
            }
        } catch (Exception e) {
            // 如果 classpath 不可用，尝试文件系统路径
        }
        
        // 如果 classpath 不可用，尝试常见的文件系统路径
        if (docsPath == null) {
            String[] possiblePaths = {
                "src/main/resources/docs",
                "classpath:docs",
                "./src/main/resources/docs"
            };
            for (String path : possiblePaths) {
                Path testPath = Paths.get(path);
                if (Files.exists(testPath) && Files.isDirectory(testPath)) {
                    docsPath = path;
                    break;
                }
            }
        }
        
        // 如果找不到 docs 目录，返回一个空的内容检索器
        if (docsPath == null || !Files.exists(Paths.get(docsPath))) {
            // 创建一个空的检索器，避免启动失败
            return EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(qwenEmbeddingModel)
                    .maxResults(5)
                    .minScore(0.7)
                    .build();
        }
        
        List<Document> documents = FileSystemDocumentLoader.loadDocuments(docsPath);
        // 2. 文档切割：将每个文档按每段进行分割，最大 1000 字符，每次重叠最多 200 个字符
        DocumentByParagraphSplitter paragraphSplitter = new DocumentByParagraphSplitter(1000, 200);
        // 3. 自定义文档加载器
        EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder()
                .textSegmentTransformer(textSegment -> TextSegment.from(
                        textSegment.metadata().getString("file_name") + "\n" + textSegment.text(),
                        textSegment.metadata()))
                // 使用指定的向量模型
                .embeddingModel(qwenEmbeddingModel)
                .documentSplitter(paragraphSplitter)
                .embeddingStore(embeddingStore)
                .build();
        // 4. 索引文档
        embeddingStoreIngestor.ingest(documents);
        // 5. 创建内容检索器
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(5) // 最多返回 5 个结果
                .minScore(0.7) // 最小分数
                .build();
        return contentRetriever;
    }
}
